# -*- coding: utf-8 -*-
# coding=utf-8

import copy
import re
from utils.log import logger
from transformers import AutoTokenizer, AutoModel
import json

class NormRecognition:
    def __init__(self, model, headers_path) -> None:
        self.model = model
        with open(headers_path, 'r') as file:
            data = json.load(file)
        self.candidates = data
        self.prompt = f"维度列表：[{'，'.join(self.candidates)}] 你将扮演一个文本规范化器，\
        将输入的文本输出为规范化的文本。输入的语句可能存在同义词、缩略语、语义缺失或隐含背景知识的问题，\
        若输入文本不存在上述问题，则输出原句，反之则输出规范化处理后的文本结果。回答内容只包含规范化的文本。\
        要求不改变原文本的语义，规范化文本不改变原输入文本中出现在上述维度列表中的词汇。"
        self.prompt = self.prompt + "文本：{}\n"

    def recognize(self, query):
        """
        识别query对应的header
        :param query: 自然语言查询
        :return: 规范化后的query
        """
        try:
            prediction = self.model.predict(self.prompt.format(query))
        except Exception as err:
            logger.error("header prediction error happened![{}]: {}".format(type(err), str(err)))
            raise
        return prediction
        

